import torch
import numpy as np
from matplotlib import pyplot as plt

x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
y = 2 * x + 1


class LRModel(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super(LRModel, self).__init__()
        self.linear = torch.nn.Linear(in_features, out_features)

    def forward(self, inputs):
        return self.linear(inputs)


model = LRModel(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.25)
loss_func = torch.nn.MSELoss()

plt.ion()
plt.show()

for time in range(100):
    optimizer.zero_grad()
    outputs = model(x)
    loss = loss_func(outputs, y)
    loss.backward()
    optimizer.step()
    if time % 5 == 0:
        plt.cla()
        plt.scatter(x.detach().numpy(), y.detach().numpy())
        plt.plot(x.detach().numpy(), outputs.detach().numpy())
        plt.text(0.5, 0, f'Loss={loss.detach()}')
        plt.pause(0.1)

plt.ioff()
plt.show()
